package com.ky.pro.java.mq.kmq.core;

import com.ky.pro.java.mq.kmq.repository.KmqMessageRepository;
import lombok.extern.slf4j.Slf4j;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

@Slf4j
public class KQueue<T> {

    long capacity;

    KmqMessageRepository repository;

    String topic;

    public KQueue(int capacity, String topic,
                  KmqMessageRepository repository) {
        this.capacity = capacity;
        this.topic = topic;
        this.repository = repository;
        if (!repository.initQueue(topic, capacity)) {
            throw new RuntimeException("初始化队列仓储失败: topic = " + topic + ", capacity = " + capacity);
        }
    }

    /**
     * 每次从队列仓储中拉取的最大消息数
     * 后续可改为在配置文件中实现
     */
    private final static int BATCH_QUERY_SIZE = 1000;

    /**
     * 每次拉取间隙的等待时间
     * 后续可改为在配置文件中实现
     */
    private final static int SLEEPING_TIME_MILLIS_BETWEEN_POLLING = 100;


    /**
     * 拉取消息
     *
     * @param consumerId
     * @param timeout
     * @param timeoutUnit
     * @param batchSize
     * @return
     */
    @Transactional
    public List<KmqMessage<T>> poll(String consumerId, long timeout, TimeUnit timeoutUnit, int batchSize) throws InterruptedException {

        long startTime = System.currentTimeMillis();

        int currentOffset = repository.getOffset(topic, consumerId);

        long elapsedTime = System.currentTimeMillis() - startTime;
        long timeoutDurationMillis = timeoutUnit.toMillis(timeout);

        int batchQuerySize = BATCH_QUERY_SIZE > batchSize ? batchSize : BATCH_QUERY_SIZE;

        List<KmqMessage<T>> msgList = new ArrayList<>();

        int consumedSize = 0;

        while (consumedSize < batchSize && elapsedTime < timeoutDurationMillis) {

            List<KmqMessage<T>> msgBatch = repository.getMessage(topic, consumerId, currentOffset,
                    batchQuerySize);
            if (msgBatch != null && !msgBatch.isEmpty()) {
                int currentBatchSize = msgBatch.size();
                msgList.addAll(msgBatch);

                currentOffset += currentBatchSize;
                consumedSize += currentBatchSize;
                if (!repository.saveOffset(topic, consumerId, currentOffset)) {
                    log.error("存储当前消费位置失败！topic = {}, consumerId = {}, currentOffset = {}",
                            topic, consumerId, currentOffset);
                    throw new RuntimeException("存储当前消费位置失败！");
                }
            }

            Thread.sleep(SLEEPING_TIME_MILLIS_BETWEEN_POLLING);
            elapsedTime = System.currentTimeMillis() - startTime;
        }


        return msgList;
    }

    /**
     * 添加消息，消息添加后当前位置
     *
     * @param msg
     * @return
     */
    long offer(KmqMessage msg) {
        return repository.saveMessageToTopic(topic, msg);
    }

    /**
     * 添加消费者
     *
     * @param consumerId
     * @return
     */
    boolean addConsumer(String consumerId) {
        return repository.addConsumerToTopic(topic, consumerId);
    }

    /**
     * 移除消费者
     */
    boolean removeConsumer(String consumerId) {
        return repository.removeConsumerFromTopic(topic, consumerId);
    }

}
